# -*- coding: utf-8 -*-
# @Time    : 2020/6/17 下午11:03
# @Author  : caotian
# @FileName: regular.py
# @Software: PyCharm
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import json
import gzip
import os
import sys
import random
import numpy as np
from PIL import Image
curpath=os.path.abspath(os.curdir)
sys.path.append(curpath)
import optimizationdata as od
import optimizationmodel as om
with fluid.dygraph.guard():
    model = om.MNIST()
    model.train()
    train_loader = od.load_data('train')
    # 四种优化算法的设置方案，可以逐一尝试效果
    optimizer = fluid.optimizer.SGDOptimizer(learning_rate=0.01, parameter_list=model.parameters())
    # optimizer = fluid.optimizer.MomentumOptimizer(learning_rate=0.01, momentum=0.9, parameter_list=model.parameters())
    # optimizer = fluid.optimizer.AdagradOptimizer(learning_rate=0.01, parameter_list=model.parameters())
    # optimizer = fluid.optimizer.AdamOptimizer(learning_rate=0.01, parameter_list=model.parameters())

    # 各种优化算法均可以加入正则化项，避免过拟合，参数regularization_coeff调节正则化项的权重
    # optimizer = fluid.optimizer.SGDOptimizer(learning_rate=0.01, regularization=fluid.regularizer.L2Decay(regularization_coeff=0.1),parameter_list=model.parameters()))
    optimizer = fluid.optimizer.AdamOptimizer(learning_rate=0.01,
                                              regularization=fluid.regularizer.L2Decay(regularization_coeff=0.1),
                                              parameter_list=model.parameters())

    EPOCH_NUM = 10
    for epoch_id in range(EPOCH_NUM):
        for batch_id, data in enumerate(train_loader()):
            # 准备数据，变得更加简洁
            image_data, label_data = data
            image = fluid.dygraph.to_variable(image_data)
            label = fluid.dygraph.to_variable(label_data)

            # 前向计算的过程，同时拿到模型输出值和分类准确率
            predict, acc = model(image, label)

            # 计算损失，取一个批次样本损失的平均值
            loss = fluid.layers.cross_entropy(predict, label)
            avg_loss = fluid.layers.mean(loss)

            # 每训练了100批次的数据，打印下当前Loss的情况
            if batch_id % 100 == 0:
                print("epoch: {}, batch: {}, loss is: {}, acc is {}".format(epoch_id, batch_id, avg_loss.numpy(),
                                                                            acc.numpy()))

            # 后向传播，更新参数的过程
            avg_loss.backward()
            optimizer.minimize(avg_loss)
            model.clear_gradients()

    # 保存模型参数
    fluid.save_dygraph(model.state_dict(), 'mnist-model')